from typing import Optional
from typing import Union
from fastapi import FastAPI
from fastapi import status
from pydantic import BaseModel
from pyvirtualdisplay import Display
from selenium import webdriver
from selenium.webdriver.chrome.options import Options
import time
from selenium.webdriver.support.ui import WebDriverWait
from selenium.webdriver.support import expected_conditions as EC
from selenium.webdriver.common.by import By
from fastapi.responses import JSONResponse, Response
import mysqlserver

app = FastAPI()


class Req(BaseModel):
    url: str
    urlcontent: str


@app.post("/")
def read_root():
    return {"Hello": "World"}


@app.get("/items/{item_id}")
def read_item(item_id: int, q: Optional[str] = None):
    return {"item_id": item_id, "q": q}



@app.post("/get_love_story")
def get_love_story():
    return resp_200(data={
        "story": mysqlserver.get_one(),
    })


@app.post("/get_html")
def get_html(req: Req):
    chrome_options = Options()
    chrome_options.add_argument('--no-sandbox')
    chrome_options.add_argument('--disable-dev-shm-usage')
    print("recive request "+req.url)
    display = Display(visible=0, size=(800, 800))
    display.start()
    try:
        driver = webdriver.Chrome(
            '/usr/local/bin/chromedriver', chrome_options=chrome_options)
        driver.get(req.url)
        driver.get(req.urlcontent)
        article = WebDriverWait(driver, 2).until(
            EC.presence_of_element_located(
                (By.CSS_SELECTOR, ".rich_media_area_primary"))
        )
        imgs = article.find_elements_by_tag_name("div img")
        list = []
        for img in imgs:
            try:
                src = img.get_attribute('data-src')
                time.sleep(0.5)
                if(src):
                    list.append(src)
            except:
                print("continue")

        return resp_200(data={
            "imgs": list,
            "text": article.text.replace("\n","\\n")
        })
    finally:
        driver.close()
        driver.quit()
        display.stop()


def resp_200(*, data: Union[list, dict, str]) -> Response:
    return JSONResponse(
        status_code=status.HTTP_200_OK,
        content=data
    )
